"""
雪花ID生成器模块
基于Twitter的Snowflake算法实现
生成64位的分布式唯一ID
"""
import time
import threading


class SnowflakeIDGenerator:
    """
    雪花ID生成器
    
    位数分配:
    - 1位: 符号位，始终为0
    - 41位: 时间戳，毫秒级
    - 10位: 机器ID，可配置
    - 12位: 序列号，同一毫秒内的自增ID
    """
    def __init__(self, worker_id=1, datacenter_id=1):
        """
        初始化雪花ID生成器
        
        Args:
            worker_id (int): 工作机器ID (0-31)
            datacenter_id (int): 数据中心ID (0-31)
        """
        # 开始时间戳 (2023-01-01 00:00:00) - 自定义起始时间
        self.start_timestamp = 1672531200000
        
        # 位数分配
        self.worker_id_bits = 5
        self.datacenter_id_bits = 5
        self.sequence_bits = 12
        
        # 最大值
        self.max_worker_id = -1 ^ (-1 << self.worker_id_bits)
        self.max_datacenter_id = -1 ^ (-1 << self.datacenter_id_bits)
        
        # 移位
        self.worker_id_shift = self.sequence_bits
        self.datacenter_id_shift = self.sequence_bits + self.worker_id_bits
        self.timestamp_shift = self.sequence_bits + self.worker_id_bits + self.datacenter_id_bits
        
        # 序列号掩码
        self.sequence_mask = -1 ^ (-1 << self.sequence_bits)
        
        # 检查参数
        if worker_id > self.max_worker_id or worker_id < 0:
            raise ValueError(f"worker_id 必须在 0 到 {self.max_worker_id} 之间")
        if datacenter_id > self.max_datacenter_id or datacenter_id < 0:
            raise ValueError(f"datacenter_id 必须在 0 到 {self.max_datacenter_id} 之间")
        
        # 初始化属性
        self.worker_id = worker_id
        self.datacenter_id = datacenter_id
        self.sequence = 0
        self.last_timestamp = -1
        
        # 线程锁
        self.lock = threading.Lock()
    
    def _get_timestamp(self):
        """获取当前毫秒时间戳"""
        return int(time.time() * 1000)
    
    def _wait_next_millis(self, last_timestamp):
        """
        等待下一毫秒
        
        Args:
            last_timestamp (int): 上次生成ID的时间戳
        
        Returns:
            int: 新的时间戳
        """
        timestamp = self._get_timestamp()
        while timestamp <= last_timestamp:
            timestamp = self._get_timestamp()
        return timestamp
    
    def next_id(self):
        """
        生成下一个ID
        
        Returns:
            int: 雪花算法生成的唯一ID
        """
        with self.lock:
            timestamp = self._get_timestamp()
            
            # 校验时钟回拨
            if timestamp < self.last_timestamp:
                # 时钟回拨，抛出异常或等待
                # 这里简单处理为使用上次时间戳
                timestamp = self.last_timestamp
            
            # 同一毫秒
            if timestamp == self.last_timestamp:
                self.sequence = (self.sequence + 1) & self.sequence_mask
                # 同一毫秒序列号用完
                if self.sequence == 0:
                    timestamp = self._wait_next_millis(self.last_timestamp)
            else:
                # 不同毫秒，序列号重置
                self.sequence = 0
            
            # 保存最后时间戳
            self.last_timestamp = timestamp
            
            # 组合ID
            return ((timestamp - self.start_timestamp) << self.timestamp_shift) | \
                   (self.datacenter_id << self.datacenter_id_shift) | \
                   (self.worker_id << self.worker_id_shift) | \
                   self.sequence


# 创建全局单例实例
_snowflake_generator = None

def init_snowflake(worker_id=1, datacenter_id=1):
    """初始化雪花ID生成器"""
    global _snowflake_generator
    if _snowflake_generator is None:
        _snowflake_generator = SnowflakeIDGenerator(worker_id, datacenter_id)

def generate_id():
    """生成雪花ID"""
    global _snowflake_generator
    if _snowflake_generator is None:
        init_snowflake()
    return _snowflake_generator.next_id() 